#include "MemHook.h"

#include <stdlib.h>
#include <stdio.h>
#include <dlfcn.h>

#include <iostream>

MallocFn real_malloc = NULL;
FreeFn real_free = NULL;

MemHook::MemHook()
	: initialized(false)
{
	// fprintf(stderr, "memhook construct\n");
	if (Init() < 0)
		exit(1);

	// 定义为指针，避免在Init前默认初始化
	// 支持多线程
	mutex = new std::mutex();
	allocate_ptrs = new AllocatorMap();
}

MemHook::~MemHook()
{
	if (allocate_ptrs) {
		std::cout << allocate_ptrs->size() << std::endl;
		for (auto& entry : (*allocate_ptrs)) {
			std::cout << entry.first << " " << entry.second << std::endl;
		}
		delete allocate_ptrs;
		allocate_ptrs = NULL;
	}
	delete mutex;
}

int MemHook::Init()
{
	if (initialized)
		return 0;

	real_malloc = reinterpret_cast<MallocFn>(dlsym(RTLD_NEXT, "malloc"));
	if (!real_malloc) {
		fprintf(stderr, "unable to get malloc symbol!\n");
		return -1;
	}

	real_free = reinterpret_cast<FreeFn>(dlsym(RTLD_NEXT, "free"));
	if (!real_free) {
		fprintf(stderr, "unable to get free symbol!\n");
		return -1;
	}

	// fprintf(stderr, "malloc wrapped\n");

	initialized = true;
	return 0;
}

void* MemHook::Malloc(size_t n)
{
	if (!initialized) {
		Init();
	}

	void* p = real_malloc(n);
	if (allocate_ptrs) {
		std::lock_guard<std::mutex> lock(*mutex);
		(*allocate_ptrs)[p] = n;
	}
	return p;
}

void MemHook::Free(void* p)
{
	if (!initialized) {
		Init();
	}

	real_free(p);
	if (allocate_ptrs) {
		std::lock_guard<std::mutex> lock(*mutex);
		(*allocate_ptrs).erase(p);
	}
}

MemHook hook;

extern "C"
{

	void* malloc(size_t n)
	{
		return hook.Malloc(n);
	}

	void free(void* p)
	{
		hook.Free(p);
	}
}
